"""
Core RWKV model wrapper
"""

# !!! 重要：必须在导入 RWKV 之前设置环境变量 !!!
import os
os.environ.setdefault("RWKV_V7_ON", "1")  # enable this for rwkv-7 models
os.environ.setdefault("RWKV_JIT_ON", "1")  # enable JIT compilation for better speed
# RWKV_CUDA_ON 由配置控制，不在这里设置

import copy
import types
import re
import time
import json
from typing import Optional, List, Tuple, Dict, Any
from pathlib import Path

import torch
import numpy as np

# 设置 torch 后端（参考代码风格）
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cuda.matmul.allow_tf32 = True
from rwkv.model import RWKV
from rwkv.utils import PIPELINE

try:
    from huggingface_hub import hf_hub_download, list_repo_files
    HF_AVAILABLE = True
except ImportError:
    HF_AVAILABLE = False
    hf_hub_download = None
    list_repo_files = None

try:
    import requests
    from bs4 import BeautifulSoup
    WEB_SCRAPING_AVAILABLE = True
except ImportError:
    WEB_SCRAPING_AVAILABLE = False
    requests = None
    BeautifulSoup = None

from .config import Config


class EasyRWKV:
    """Easy-to-use RWKV model wrapper"""
    
    # Popular RWKV models on Hugging Face
    HF_MODELS = {
        'rwkv-4-pile-169m': {
            'repo': 'RWKV/rwkv-4-169m-pile',
            'filename': 'RWKV-4-Pile-169M-20220807-8023.pth',
            'size': '170MB',
            'description': 'RWKV 4 Pile 169M - Small model for testing'
        },
        'rwkv-5-eagle-7b': {
            'repo': 'RWKV/v5-Eagle-7B-HF',
            'filename': 'RWKV-5-World-7B-v2-20240128-ctx4096.pth',
            'size': '7GB',
            'description': 'RWKV 5 Eagle 7B - High quality world model'
        },
        'rwkv-6-world-7b': {
            'repo': 'RWKV/rwkv-6-world-7b',
            'filename': 'RWKV-x060-World-7B-v2.1-20240507-ctx4096.pth',
            'size': '7GB',
            'description': 'RWKV 6 World 7B - Latest world model'
        }
    }
    
    def __init__(self, config: Optional[Config] = None, model_path: Optional[str] = None):
        """
        Initialize Easy RWKV
        
        Args:
            config: Configuration object
            model_path: Path to RWKV model, local path, or Hugging Face model name
        """
        self.config = config or Config()
        
        if model_path:
            # Check if it's a Hugging Face model name or local path
            if self._is_hf_model_name(model_path):
                self.config.model_path = self._download_hf_model(model_path)
            else:
                self.config.model_path = model_path
        
        if not self.config.model_path:
            raise ValueError("Model path must be provided either through config or model_path parameter")
        
        # Set environment variables
        os.environ["RWKV_V7_ON"] = "1" if self.config.rwkv_v7_on else "0"
        os.environ["RWKV_JIT_ON"] = "1" if self.config.rwkv_jit_on else "0"
        os.environ["RWKV_CUDA_ON"] = "1" if self.config.rwkv_cuda_on else "0"
        
        # Set torch backends for better performance
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.allow_tf32 = True
        torch.backends.cuda.matmul.allow_tf32 = True
        
        self.model: Optional[RWKV] = None
        self.pipeline: Optional[PIPELINE] = None
        self.model_tokens: List[int] = []
        self.model_state: Optional[List] = None
        self._is_loaded = False
        
    def load_model(self) -> None:
        """Load the RWKV model"""
        if self._is_loaded:
            return
        
        # 确保在加载模型前设置正确的环境变量
        os.environ["RWKV_V7_ON"] = "1" if self.config.rwkv_v7_on else "0"
        os.environ["RWKV_JIT_ON"] = "1" if self.config.rwkv_jit_on else "0"
        os.environ["RWKV_CUDA_ON"] = "1" if self.config.rwkv_cuda_on else "0"
        
        # Check if we need to download the model first
        model_path = self.config.model_path
        if self._is_hf_model_name(model_path):
            print(f"🤗 Detected Hugging Face model: {model_path}")
            model_path = self._download_hf_model(model_path)
            # Update config with the local path
            self.config.model_path = model_path

        print(f"Loading model - {model_path}")
        
        try:
            # 新版 RWKV 库会自动添加 .pth 扩展名
            # 如果路径已经以 .pth 结尾，我们需要去掉它
            rwkv_model_path = model_path
            if rwkv_model_path.endswith('.pth'):
                rwkv_model_path = rwkv_model_path[:-4]  # 去掉 .pth
                
            self.model = RWKV(model=rwkv_model_path, strategy=self.config.strategy)
            self.pipeline = PIPELINE(self.model, "rwkv_vocab_v20230424")
            
            # 修复 RWKV v7.0 模型兼容性问题：添加缺失的 n_head 属性
            if not hasattr(self.model.args, 'n_head'):
                # 根据权重结构推断 n_head
                # 通常 n_head = n_att // head_size，从 r_k 权重可以推断
                if hasattr(self.model.args, 'n_att'):
                    # 对于这个特定模型，我们知道 n_head=12，head_size=64
                    head_size = 64  # 从权重分析得出
                    n_head = self.model.args.n_att // head_size
                    self.model.args.n_head = n_head
                    print(f"💡 添加缺失的 n_head 属性: {n_head}")
                else:
                    # 如果连 n_att 都没有，使用默认值
                    self.model.args.n_head = 12
                    print("💡 使用默认 n_head 值: 12")
            
            # Load custom state if provided
            if self.config.state_path and os.path.exists(self.config.state_path):
                self._load_custom_state()
            
            self._is_loaded = True
            print("Model loaded successfully!")
            
        except Exception as e:
            error_msg = str(e).lower()
            
            # 检测CUDA相关错误并提供友好的错误信息
            if "cuda" in error_msg or "torch not compiled with cuda" in error_msg:
                print("\n" + "="*60)
                print("🚨 CUDA 配置问题")
                print("="*60)
                print(f"❌ 错误信息: {e}")
                print("\n💡 解决方案:")
                print("1. 在 macOS 上请使用 CPU 策略:")
                print("   --strategy 'cpu fp32'")
                print("\n2. 或者更新配置:")
                print("   config.strategy = 'cpu fp32'")
                print("   config.rwkv_cuda_on = False")
                print("\n3. 确保环境变量正确:")
                print("   export RWKV_CUDA_ON=0")
                print("\n🔧 推荐的完整命令:")
                print(f"   easy-rwkv chat --model '{self.config.model_path}' --strategy 'cpu fp32'")
                print("="*60)
                raise RuntimeError("CUDA配置错误。请参考上述解决方案使用CPU模式。")
            
            # 检测模型文件相关错误
            elif "file not found" in error_msg or "no such file" in error_msg:
                print("\n" + "="*60)
                print("🚨 模型文件问题")
                print("="*60)
                print(f"❌ 错误信息: {e}")
                print("\n💡 解决方案:")
                print("1. 检查模型路径是否正确")
                print("2. 如果是HuggingFace模型，请先下载:")
                print("   easy-rwkv models  # 查看可用模型")
                print("3. 或者使用完整路径:")
                print("   /Users/tangyubin/.cache/easy_rwkv/models/your-model.pth")
                print("="*60)
                raise RuntimeError("模型文件未找到。请检查路径或先下载模型。")
            
            else:
                # 其他未知错误，保持原有行为但增加一些提示
                print(f"\n❌ 模型加载失败: {e}")
                print("\n💡 常见解决方案:")
                print("1. 确保使用正确的策略: --strategy 'cpu fp32'")
                print("2. 检查模型文件是否存在")
                print("3. 查看完整错误信息并报告给开发者")
                raise RuntimeError(f"Failed to load model: {e}")
    
    def _load_custom_state(self) -> None:
        """Load custom state from file"""
        print(f"Loading custom state - {self.config.state_path}")
        
        args = self.model.args
        state_raw = torch.load(self.config.state_path + '.pth')
        state_init = [None for i in range(args.n_layer * 3)]
        
        for i in range(args.n_layer):
            dd = self.model.strategy[i]
            dev = dd.device
            atype = dd.atype
            
            state_init[i*3+0] = torch.zeros(args.n_embd, dtype=atype, requires_grad=False, device=dev).contiguous()
            state_init[i*3+1] = state_raw[f'blocks.{i}.att.time_state'].transpose(1,2).to(dtype=torch.float, device=dev).requires_grad_(False).contiguous()
            state_init[i*3+2] = torch.zeros(args.n_embd, dtype=atype, requires_grad=False, device=dev).contiguous()
        
        self.model_state = copy.deepcopy(state_init)
    
    def _run_rnn(self, ctx: str) -> torch.Tensor:
        """Run RNN forward pass"""
        if not self._is_loaded:
            self.load_model()
        
        ctx = ctx.replace("\r\n", "\n")
        
        tokens = self.pipeline.encode(ctx)
        tokens = [int(x) for x in tokens]
        self.model_tokens += tokens
        
        while len(tokens) > 0:
            # 使用参考代码的方式处理状态
            if self.model_state is None:
                # 第一次调用，让 RWKV 自动初始化状态
                out, self.model_state = self.model.forward(tokens[:self.config.chunk_len], None)
            else:
                out, self.model_state = self.model.forward(tokens[:self.config.chunk_len], self.model_state)
            tokens = tokens[self.config.chunk_len:]
        
        return out
    
    def initialize_chat(self) -> None:
        """Initialize chat with system prompt"""
        if not self._is_loaded:
            self.load_model()
        
        if self.config.state_path is None:
            self._run_rnn(self.config.system_prompt)
    
    def chat(self, message: str, max_tokens: Optional[int] = None, temperature: Optional[float] = None, 
             top_p: Optional[float] = None) -> str:
        """
        Chat with the model
        
        Args:
            message: User input message
            max_tokens: Maximum tokens to generate (overrides config)
            temperature: Temperature for sampling (overrides config)
            top_p: Top-p for sampling (overrides config)
        
        Returns:
            Model response
        """
        if not self._is_loaded:
            self.load_model()
        
        # Use provided parameters or fall back to config
        max_tokens = max_tokens or self.config.max_tokens
        temperature = temperature if temperature is not None else self.config.temperature
        top_p = top_p if top_p is not None else self.config.top_p
        
        # Clean up input message
        message = message.strip()
        message = re.sub(r"\n+", "\n", message)
        
        if len(message) == 0:
            raise ValueError("Message cannot be empty")
        
        occurrence = {}
        out_tokens = []
        out_last = 0
        response_parts = []
        
        # Run the input through the model
        out = self._run_rnn("User: " + message + "\n\nAssistant:")
        
        # Generate response
        for i in range(max_tokens):
            # Apply repetition penalty
            for n in occurrence:
                out[n] -= self.config.alpha_presence + occurrence[n] * self.config.alpha_frequency
            
            out[0] -= 1e10  # disable END_OF_TEXT
            
            # Sample next token
            token = self.pipeline.sample_logits(out, temperature=temperature, top_p=top_p)
            
            # Forward pass
            out, self.model_state = self.model.forward([token], self.model_state)
            self.model_tokens += [token]
            out_tokens += [token]
            
            # Update occurrence tracking
            for xxx in occurrence:
                occurrence[xxx] *= self.config.penalty_decay
            occurrence[token] = 1 + (occurrence[token] if token in occurrence else 0)
            
            # Decode and check for completion
            tmp = self.pipeline.decode(out_tokens[out_last:])
            
            if ("\ufffd" not in tmp) and (not tmp.endswith("\n")):
                response_parts.append(tmp)
                out_last = i + 1
            
            if "\n\n" in tmp:
                response_parts.append(tmp)
                break
        
        return "".join(response_parts).strip()
    
    def reset_state(self) -> None:
        """Reset model state"""
        self.model_tokens = []
        self.model_state = None
        if self.config.state_path and os.path.exists(self.config.state_path):
            self._load_custom_state()
    
    def save_state(self, state_path: str) -> None:
        """Save current model state"""
        if self.model_state is not None:
            torch.save(self.model_state, state_path)
        else:
            raise RuntimeError("No state to save. Model might not be loaded or initialized.")
    
    def get_model_info(self) -> Dict[str, Any]:
        """Get model information"""
        if not self._is_loaded:
            return {"status": "not_loaded"}
        
        return {
            "status": "loaded",
            "model_path": self.config.model_path,
            "strategy": self.config.strategy,
            "state_path": self.config.state_path,
            "tokens_processed": len(self.model_tokens),
            "has_state": self.model_state is not None,
        }
    
    @staticmethod
    def _is_hf_model_name(model_path: str) -> bool:
        """Check if model_path is a Hugging Face model name rather than a local path"""
        # It's likely a HF model name if:
        # 1. It's in our known models list
        # 2. It contains "/" (repo format) and doesn't exist as a local file
        # 3. It's not an absolute path (starts with / or \)
        
        if model_path in EasyRWKV.HF_MODELS:
            return True
        
        # Check for absolute paths
        if model_path.startswith('/') or model_path.startswith('\\') or ':' in model_path[:3]:
            return False
            
        # If it contains "/" and doesn't exist as a local file, treat as HF model
        if "/" in model_path and not Path(model_path).exists():
            return True
            
        # If it doesn't exist as a local file and contains no path separators
        if not Path(model_path).exists() and '/' not in model_path and '\\' not in model_path:
            return False
            
        return False
    
    def _download_hf_model(self, model_name: str) -> str:
        """Automatically download RWKV model from Hugging Face with progress bar"""
        
        print(f"🤗 检测到 Hugging Face 模型: {model_name}")
        
        # Check if it's a known model shortname
        if model_name in self.HF_MODELS:
            model_info = self.HF_MODELS[model_name]
            repo_id = model_info['repo']
            filename = model_info['filename']
            print(f"📦 预定义模型: {model_info['description']}")
            print(f"📏 模型大小: {model_info['size']}")
        else:
            # Assume it's a repo_id/filename format or just repo_id
            if "/" in model_name:
                if ".pth" in model_name:
                    # It's a full path: repo_id/filename
                    parts = model_name.split("/")
                    repo_id = "/".join(parts[:-1])
                    filename = parts[-1]
                    print(f"📁 指定文件: {filename} 来自 {repo_id}")
                else:
                    # Just repo_id, need to find .pth files
                    repo_id = model_name
                    try:
                        # Try HF API first
                        if HF_AVAILABLE and list_repo_files is not None:
                            files = list_repo_files(repo_id)
                            pth_files = [f for f in files if f.endswith('.pth')]
                            if pth_files:
                                filename = pth_files[0]  # Use the first .pth file
                                print(f"📁 在仓库中找到 {len(pth_files)} 个 .pth 文件，将使用: {filename}")
                            else:
                                raise ValueError("仓库中未找到 .pth 文件")
                        else:
                            raise ImportError("huggingface_hub 不可用于文件列表")
                    except Exception as hf_error:
                        # Fallback to web scraping
                        print(f"⚠️  HF API 失败 ({hf_error})，尝试网页抓取...")
                        pth_files = self.__class__._get_pth_files_from_repo(repo_id)
                        if not pth_files:
                            raise ValueError(f"在仓库 {repo_id} 中未找到 .pth 文件")
                        filename = pth_files[0]  # Use the first .pth file
                        print(f"📁 通过网页抓取找到 {len(pth_files)} 个 .pth 文件，将使用: {filename}")
            else:
                raise ValueError(f"未知的模型格式: {model_name}")
        
        # Create cache directory
        cache_dir = Path.home() / ".cache" / "easy_rwkv" / "models"
        cache_dir.mkdir(parents=True, exist_ok=True)
        
        # Check if model already exists in our simple cache
        possible_paths = list(cache_dir.rglob(filename))
        if possible_paths:
            local_path = str(possible_paths[0])
            print(f"✅ 模型已存在于简单缓存中: {local_path}")
            return local_path
        
        # Check if model exists in HuggingFace cache
        if HF_AVAILABLE:
            try:
                from huggingface_hub import hf_hub_download
                # 尝试获取缓存中的文件路径
                hf_cached_path = hf_hub_download(
                    repo_id=repo_id,
                    filename=filename,
                    cache_dir=str(cache_dir),
                    local_files_only=True  # 只检查缓存，不下载
                )
                if hf_cached_path and os.path.exists(hf_cached_path):
                    print(f"✅ 模型已存在于HF缓存中: {hf_cached_path}")
                    return hf_cached_path
            except Exception:
                # 缓存中没有文件，继续下载流程
                pass
        
        # Model not found, try automatic download
        if not HF_AVAILABLE:
            # Fallback to manual download instructions
            return self._show_manual_download_instructions(repo_id, filename, cache_dir)
        
        try:
            print(f"\n🚀 开始自动下载模型文件...")
            print(f"📁 仓库: {repo_id}")
            print(f"📄 文件: {filename}")
            print(f"💾 缓存目录: {cache_dir}")
            
            # Import necessary modules for progress tracking
            from huggingface_hub import hf_hub_download
            
            # Try to import tqdm for progress bar
            try:
                from tqdm.auto import tqdm
                TQDM_AVAILABLE = True
            except ImportError:
                TQDM_AVAILABLE = False
                print("📊 (安装 tqdm 可获得详细进度条: pip install tqdm)")
            
            # Download with progress tracking
            print(f"⏬ 正在下载，请稍等...")
            if TQDM_AVAILABLE:
                print("📈 显示下载进度...")
            
            local_path = hf_hub_download(
                repo_id=repo_id,
                filename=filename,
                cache_dir=str(cache_dir),
                resume_download=True,
                # Let huggingface_hub handle the progress bar automatically
            )
            
            print(f"\n✅ 模型下载成功!")
            print(f"📍 本地路径: {local_path}")
            return local_path
            
        except Exception as e:
            error_str = str(e)
            
            # Handle authentication errors
            if "401" in error_str or "Unauthorized" in error_str or "authentication" in error_str.lower():
                print(f"\n🔐 认证错误：{e}")
                print("\n💡 解决方案:")
                print("1. 登录 Hugging Face:")
                print("   huggingface-cli login")
                print("   # 然后输入你的 HF token (从 https://huggingface.co/settings/tokens 获取)")
                print("\n2. 或设置 HF_TOKEN 环境变量:")
                print("   export HF_TOKEN=your_token_here")
                print("\n3. 某些模型可能需要申请访问权限")
                print(f"   请访问 https://huggingface.co/{repo_id} 查看是否需要申请权限")
                
                # Fall back to manual instructions
                return self._show_manual_download_instructions(repo_id, filename, cache_dir)
            
            # Handle network errors
            elif "connection" in error_str.lower() or "timeout" in error_str.lower():
                print(f"\n🌐 网络错误：{e}")
                print("\n💡 解决方案:")
                print("1. 检查网络连接")
                print("2. 尝试使用镜像站点:")
                print("   export HF_ENDPOINT=https://hf-mirror.com")
                print("3. 重试下载")
                
                # Fall back to manual instructions
                return self._show_manual_download_instructions(repo_id, filename, cache_dir)
            
            # Handle file not found errors
            elif "404" in error_str or "not found" in error_str.lower():
                print(f"\n❌ 文件未找到：{e}")
                print(f"\n💡 可能的原因:")
                print(f"1. 仓库 {repo_id} 不存在")
                print(f"2. 文件 {filename} 不存在")
                print("3. 仓库是私有的，需要认证")
                print("\n🔍 建议:")
                print(f"   访问 https://huggingface.co/{repo_id} 确认文件存在")
                
                # Fall back to manual instructions
                return self._show_manual_download_instructions(repo_id, filename, cache_dir)
            
            # Other errors
            else:
                print(f"\n❌ 下载失败：{e}")
                print("\n🔄 回退到手动下载指导...")
                
                # Fall back to manual instructions
                return self._show_manual_download_instructions(repo_id, filename, cache_dir)
    
    def _show_manual_download_instructions(self, repo_id: str, filename: str, cache_dir: Path) -> str:
        """Show manual download instructions and raise error"""
        print("\n" + "="*70)
        print("🚨 模型文件未找到，需要手动下载")
        print("="*70)
        
        print(f"\n📋 请按以下步骤下载模型:")
        print(f"📁 目标目录: {cache_dir}")
        print(f"📄 文件名: {filename}")
        
        # Check for HF_ENDPOINT environment variable
        hf_endpoint = os.environ.get('HF_ENDPOINT', 'https://huggingface.co')
        if hf_endpoint != 'https://huggingface.co':
            print(f"🌐 当前使用镜像: {hf_endpoint}")
        
        print(f"\n🔽 方式一：使用 huggingface-cli 下载（推荐）")
        print(f"   huggingface-cli download {repo_id} {filename} --cache-dir \"{cache_dir}\"")
        
        print(f"\n🔽 方式二：手动下载")
        print(f"   # 访问: {hf_endpoint}/{repo_id}")
        print(f"   # 下载文件保存到: {cache_dir}/{filename}")
        
        print(f"\n📝 下载完成后，请重新运行您的程序")
        print("="*70)
        
        raise RuntimeError(f"模型文件 {filename} 未找到。请按照上述说明下载模型后重新运行程序。")
    
    def _download_hf_model_fallback(self, model_name: str) -> str:
        """Fallback download method using huggingface_hub library"""
        if not HF_AVAILABLE:
            raise ImportError(
                "huggingface_hub is required for downloading models. "
                "Install it with: pip install huggingface_hub"
            )
        
        print(f"📚 Using library method to download: {model_name}")
        
        # Check if it's a known model shortname
        if model_name in self.HF_MODELS:
            model_info = self.HF_MODELS[model_name]
            repo_id = model_info['repo']
            filename = model_info['filename']
            print(f"📦 Using predefined model: {model_info['description']}")
            print(f"📏 Model size: {model_info['size']}")
        else:
            # Assume it's a repo_id/filename format or just repo_id
            if "/" in model_name:
                if ".pth" in model_name:
                    # It's a full path: repo_id/filename
                    parts = model_name.split("/")
                    repo_id = "/".join(parts[:-1])
                    filename = parts[-1]
                else:
                    repo_id = model_name
                    # Try to find .pth files in the repo
                    try:
                        # First try HF API
                        if list_repo_files is not None:
                            files = list_repo_files(repo_id)
                            pth_files = [f for f in files if f.endswith('.pth')]
                            if pth_files:
                                filename = pth_files[0]  # Use the first .pth file
                                print(f"📁 Found {len(pth_files)} .pth file(s) via HF API, using: {filename}")
                            else:
                                raise ValueError("No .pth files found via HF API")
                        else:
                            raise ImportError("huggingface_hub not available")
                    except Exception as hf_error:
                        # Fallback to web scraping
                        print(f"⚠️  HF API failed ({hf_error}), trying web scraping...")
                        pth_files = self.__class__._get_pth_files_from_repo(repo_id)
                        if not pth_files:
                            raise ValueError(f"No .pth files found in repository {repo_id}")
                        filename = pth_files[0]  # Use the first .pth file
                        print(f"📁 Found {len(pth_files)} .pth file(s) via web scraping, using: {filename}")
            else:
                raise ValueError(f"Unknown model format: {model_name}")
        
        try:
            # Create cache directory
            cache_dir = Path.home() / ".cache" / "easy_rwkv" / "models"
            cache_dir.mkdir(parents=True, exist_ok=True)
            
            print(f"📥 开始下载模型文件...")
            print(f"🔗 从 {repo_id} 下载 {filename}")
            print(f"📁 缓存目录: {cache_dir}")
            print(f"⏳ 正在下载，请稍等...")
            
            if hf_hub_download is not None:
                # Use tqdm callback for progress if available
                try:
                    from tqdm.auto import tqdm
                    
                    class ProgressCallback:
                        def __init__(self):
                            self.pbar = None
                            
                        def __call__(self, downloaded_bytes: int, total_bytes: int = None):
                            if self.pbar is None and total_bytes:
                                self.pbar = tqdm(
                                    total=total_bytes,
                                    unit='B',
                                    unit_scale=True,
                                    desc="下载进度"
                                )
                            if self.pbar:
                                self.pbar.update(downloaded_bytes - self.pbar.n)
                                
                        def close(self):
                            if self.pbar:
                                self.pbar.close()
                    
                    progress_callback = ProgressCallback()
                    local_path = hf_hub_download(
                        repo_id=repo_id,
                        filename=filename,
                        cache_dir=str(cache_dir),
                        resume_download=True
                    )
                    progress_callback.close()
                    
                except ImportError:
                    # No tqdm available, download without progress bar
                    print("📊 (无法显示详细进度，tqdm 未安装)")
                    local_path = hf_hub_download(
                        repo_id=repo_id,
                        filename=filename,
                        cache_dir=str(cache_dir),
                        resume_download=True
                    )
            else:
                raise ImportError("huggingface_hub not available")
            
            print(f"✅ Model downloaded successfully!")
            print(f"📍 Local path: {local_path}")
            
            return local_path
            
        except Exception as e:
            error_str = str(e)
            if "401" in error_str and "Unauthorized" in error_str:
                print("\n🔐 认证错误！BlinkDL 的模型需要 Hugging Face 登录才能下载。")
                print("\n💡 解决方案:")
                print("1. 登录 Hugging Face:")
                print("   huggingface-cli login")
                print("   # 然后输入你的 HF token (从 https://huggingface.co/settings/tokens 获取)")
                print("\n2. 或者手动下载模型文件到本地:")
                print(f"   # 从 https://huggingface.co/{repo_id} 下载 {filename}")
                print(f"   # 然后使用本地路径加载模型")
                print("\n3. 或者使用预定义的不需要认证的模型:")
                print("   # 查看可用模型: python -m easy_rwkv models")
                raise RuntimeError(f"Authentication required for model {model_name}. Please login with 'huggingface-cli login' first.")
            else:
                raise RuntimeError(f"Failed to download model {model_name}: {e}")
    
    @classmethod
    def _get_cache_path(cls) -> Path:
        """Get cache directory path"""
        cache_dir = Path.home() / ".cache" / "easy_rwkv"
        cache_dir.mkdir(parents=True, exist_ok=True)
        return cache_dir / "blinkdl_models.json"
    
    @classmethod
    def _load_cached_models(cls) -> Optional[Dict[str, Dict[str, Any]]]:
        """Load models from cache if available and not expired"""
        cache_path = cls._get_cache_path()
        if not cache_path.exists():
            return None
        
        try:
            with open(cache_path, 'r', encoding='utf-8') as f:
                cache_data = json.load(f)
            
            # Check if cache is expired (24 hours)
            cache_time = cache_data.get('timestamp', 0)
            current_time = time.time()
            if current_time - cache_time > 24 * 3600:
                return None
            
            return cache_data.get('models', {})
        except Exception:
            return None
    
    @classmethod
    def clear_models_cache(cls) -> None:
        """Clear the cached BlinkDL models"""
        cache_path = cls._get_cache_path()
        if cache_path.exists():
            cache_path.unlink()
            print("🗑️  Cleared BlinkDL models cache")
        else:
            print("📭 No cache file found")

    @classmethod
    def _save_cached_models(cls, models: Dict[str, Dict[str, Any]]) -> None:
        """Save models to cache"""
        cache_path = cls._get_cache_path()
        cache_data = {
            'timestamp': time.time(),
            'models': models
        }
        
        try:
            with open(cache_path, 'w', encoding='utf-8') as f:
                json.dump(cache_data, f, indent=2, ensure_ascii=False)
        except Exception:
            pass  # Ignore cache save errors

    @classmethod
    def get_blinkdl_models(cls) -> Dict[str, Dict[str, Any]]:
        """Get all RWKV models from BlinkDL's Hugging Face page"""
        if not WEB_SCRAPING_AVAILABLE:
            print("⚠️  Web scraping not available. Install requests and beautifulsoup4.")
            return {}
        
        # Try to load from cache first
        cached_models = cls._load_cached_models()
        if cached_models:
            print("📋 Using cached BlinkDL models (cache expires in 24 hours)")
            return cached_models
        
        try:
            # Check for HF_ENDPOINT environment variable for mirror sites
            hf_endpoint = os.environ.get('HF_ENDPOINT', 'https://huggingface.co')
            blinkdl_url = f"{hf_endpoint}/BlinkDL"
            
            print(f"🔍 Fetching RWKV models from BlinkDL's page: {blinkdl_url}")
            
            # Get the main page
            response = requests.get(blinkdl_url, timeout=30)
            response.raise_for_status()
            
            soup = BeautifulSoup(response.content, 'html.parser')
            models = {}
            
            # Find all model repository links
            repo_links = soup.find_all('a', href=re.compile(r'^/BlinkDL/[^/]+$'))
            
            print(f"📦 Found {len(repo_links)} repositories...")
            
            for link in repo_links:
                repo_name = link['href'].split('/')[-1]
                if 'rwkv' in repo_name.lower():
                    repo_id = f"BlinkDL/{repo_name}"
                    print(f"   Processing {repo_id}...")
                    
                    # Get .pth files from this repository
                    pth_files = cls._get_pth_files_from_repo(repo_id)
                    
                    if pth_files:
                        models[repo_name] = {
                            'repo': repo_id,
                            'files': pth_files,
                            'description': f'RWKV model from {repo_name}'
                        }
            
            print(f"✅ Found {len(models)} RWKV model repositories with .pth files")
            
            # Save to cache
            cls._save_cached_models(models)
            
            return models
            
        except requests.exceptions.ConnectTimeout:
            print("❌ Connection timeout. Try setting HF_ENDPOINT environment variable for mirror sites.")
            print("   Example: export HF_ENDPOINT=https://hf-mirror.com")
            return {}
        except requests.exceptions.ConnectionError:
            print("❌ Connection error. Check your network or try setting HF_ENDPOINT environment variable.")
            print("   Example: export HF_ENDPOINT=https://hf-mirror.com")
            return {}
        except Exception as e:
            print(f"❌ Error fetching BlinkDL models: {e}")
            return {}

    @classmethod
    def _get_pth_files_from_repo(cls, repo_id: str) -> List[str]:
        """Get .pth files from a specific repository"""
        # Try HF API first, then fallback to web scraping
        if HF_AVAILABLE:
            try:
                if list_repo_files is not None:
                    files = list_repo_files(repo_id)
                    pth_files = [f for f in files if f.endswith('.pth')]
                    if pth_files:
                        return pth_files
            except Exception:
                pass  # Fall back to web scraping
        
        # Fallback: web scraping for .pth files
        if not WEB_SCRAPING_AVAILABLE:
            return []
        
        try:
            hf_endpoint = os.environ.get('HF_ENDPOINT', 'https://huggingface.co')
            repo_url = f"{hf_endpoint}/{repo_id}/tree/main"
            
            response = requests.get(repo_url, timeout=30)
            response.raise_for_status()
            
            soup = BeautifulSoup(response.content, 'html.parser')
            pth_files = []
            
            # Look for .pth file links
            for link in soup.find_all('a'):
                href = link.get('href', '')
                if '.pth' in href and '/blob/' in href:
                    # Extract filename from the href
                    filename = href.split('/')[-1]
                    if filename.endswith('.pth'):
                        pth_files.append(filename)
            
            return list(set(pth_files))  # Remove duplicates
            
        except Exception:
            return []

    @classmethod
    def list_available_models(cls) -> None:
        """List available pre-defined Hugging Face models"""
        print("🤗 Available Hugging Face RWKV Models:")
        print("=" * 60)
        
        for name, info in cls.HF_MODELS.items():
            print(f"📦 {name}")
            print(f"   Description: {info['description']}")
            print(f"   Repository: {info['repo']}")
            print(f"   Size: {info['size']}")
            print()
        
    @classmethod
    def list_all_models(cls) -> None:
        """List all available models from BlinkDL's Hugging Face"""
        print("🌟 All BlinkDL RWKV Models (from Hugging Face):")
        print("=" * 60)
        
        # Get BlinkDL models
        blinkdl_models = cls.get_blinkdl_models()
        
        if not blinkdl_models:
            print("❌ Could not fetch BlinkDL models. Check your internet connection.")
            print("💡 Try setting HF_ENDPOINT environment variable for mirror sites:")
            print("   export HF_ENDPOINT=https://hf-mirror.com")
            return
        
        for repo_name, info in blinkdl_models.items():
            print(f"📦 {repo_name}")
            print(f"   Repository: {info['repo']}")
            print(f"   Description: {info['description']}")
            print(f"   Available files ({len(info['files'])}):")
            
            # Show all files instead of limiting to 5
            for file in info['files']:
                print(f"     • {file}")
            
            print(f"   💡 Usage: EasyRWKV(model_path='{info['repo']}')")
            print()
        
        print("💡 You can use any repository ID or specific file:")
        print("   rwkv = EasyRWKV(model_path='BlinkDL/rwkv-7-world')")
        print("   rwkv = EasyRWKV(model_path='BlinkDL/rwkv-7-world/RWKV-x070-World-1.5B-v3-20250127-ctx4096.pth')")
        print()
